import os, sys, subprocess, math

default_header = ["index", "name", "temperature.gpu", "power.draw", "power.limit", "memory.free", "memory.used", "memory.total", "utilization.gpu"]

def gpus_snapshot(idxs:list = None):
  if sys.platform == "win32":
    nvidia = "nvidia-smi.exe"
  else:
    nvidia = "nvidia-smi"
  query = ",".join(x for x in default_header)
  cmd = [
    nvidia, 
    "--query-gpu", query,
    "--format=csv,noheader,nounits",
  ]
  if isinstance(idxs, list):
    idxs = ",".join(str(x) for x in idxs)
    cmd += ["-i", idxs]
  res = subprocess.run(cmd, stdout=subprocess.PIPE)
  if res.returncode == 0:
    res = res.stdout.decode("utf-8").strip()
  else:
    print("Error:", res.stderr.strip())
    return []
  ret = sorted([dict(zip(default_header, x.split(", "))) for x in res.split("\n")], key=lambda x: x.get("index"))
  idxs = [int(x.get("index")) for x in ret]
  return ret, idxs

class GPU:
  def __init__(self, index, auto_update=True, **kwargs):
    self.auto_update = auto_update
    self.__index = index
    self.__states = kwargs
    
  @property
  def index(self):
    return self.__index
  
  @property
  def name(self):
    if self.auto_update:
      self.update(["name"])
    return self.__states.get("name", "Unknown")
  
  @property
  def temperature(self):
    if self.auto_update:
      self.update(["temperature.gpu"])
    return float(self.__states.get("temperature.gpu", math.nan))
  
  @property
  def power(self):
    if self.auto_update:
      self.update(["power.draw"])
    return float(self.__states.get("power.draw", math.nan))
  
  @property
  def power_limit(self):
    if self.auto_update:
      self.update(["power.limit"])
    return float(self.__states.get("power.limit", math.nan))
  
  @property
  def memory_used(self):
    if self.auto_update:
      self.update(["memory.used"])
    return float(self.__states.get("memory.used", math.nan))
  
  @property
  def memory_total(self):
    if self.auto_update:
      self.update(["memory.total"])
    return float(self.__states.get("memory.total", math.nan))
  
  @property
  def memory_free(self):
    if self.auto_update:
      self.update(["memory.free"])
    return float(self.__states.get("memory.free", math.nan))
  
  @property
  def load(self):
    if self.auto_update:
      self.update(["utilization.gpu"])
    return float(self.__states.get("utilization.gpu", math.nan)) / 100
  
  @property
  def memory_percent(self):
    total = self.memory_total
    used = self.memory_used
    if math.isnan(total) or math.isnan(used):
      return math.nan
    return float(used) / float(total)
  
  @property
  def power_percent(self):
    limit = self.power_limit
    draw = self.power
    if math.isnan(limit) or math.isnan(draw):
      return math.nan
    return float(draw) / float(limit)

  def update_from_dict(self, d):
    if d.get("index") != self.__index:
      print("GPU index mismatch", self.__index, d.get("index"))
      return
    self.__states.update(d)
    
  def update(self, header=default_header):
    if sys.platform == "win32":
      nvidia = "nvidia-smi.exe"
    else:
      nvidia = "nvidia-smi"
    query = ",".join(x for x in header)
    idx = self.__index
    res = subprocess.run([
      nvidia, 
      "-i", str(idx), 
      "--query-gpu", query,
      "--format", "csv,noheader,nounits",
      ],
      stdout=subprocess.PIPE
      ).stdout.decode("utf-8").strip()
    val = res.split(", ")
    ret = dict(zip(header, val))
    self.__states.update(ret)
    
  @property
  def format(self):
    strs = f"{self.index:>3}  {self.temperature:>4}  {self.power_percent*100:>4.2f}  {self.load*100:>4.2f}  {self.memory_percent*100:>4.2f}"
    return strs
  
  def __repr__(self):
    strs = f"GPU:{self.index} at {hex(id(self))}"
    return strs